#include <Windows.h>
#include "KinectRender.qt.h"
#include "Utils.h"
#include <string>
using namespace std;

#define  LINEWIDTH          6.0f
#define  HEADDIAMETER       50.0f
#define  HANDFOOTDIAMETER   24.0f
#define  JOINTDIAMETER      14.0f

// class KinectQtRender
static void verifyBitmap(int srcWidth, int srcHeight, QImage::Format format, QImage *& dstBitmap)
{
    if(dstBitmap != NULL)
    {
        if((dstBitmap->width() != srcWidth) || (dstBitmap->height() != srcHeight))
        {
            delete dstBitmap;
            dstBitmap = NULL;
        }
    }
    if(dstBitmap == NULL)
    {
        dstBitmap = new QImage(srcWidth, srcHeight, format);
    }
}

static int copyBitsToBitmap(const void *srcBits, int srcStride, QImage *dstBitmap)
{
    const int width  = dstBitmap->width();
    const int height = dstBitmap->height();
    const int dstStride = dstBitmap->bytesPerLine();
    if(dstStride == srcStride)
    {
        memcpy(dstBitmap->bits(), srcBits, srcStride*height);
    }
    else
    {
        byte *src = (byte *)(srcBits);
        byte *dst = (byte *)(dstBitmap->bits());
        int stride = (srcStride < dstStride) ? srcStride : dstStride;
        for(int h = 0; h < height; h ++)
        {
            memcpy(dst, src, stride);
            src += srcStride;
            dst += dstStride;
        }
    }
    return 0;
}

static void drawBitmapToBitmap(QImage *dstBitmap, QImage *srcBitmap)
{
    QPainter painter(dstBitmap);
    painter.setRenderHint(QPainter::RenderHint::SmoothPixmapTransform, true);
    painter.drawImage(QRect(0, 0, dstBitmap->width(), dstBitmap->height()), *srcBitmap);
}

bool KinectQtRender::saveImageToFile(QImage *image, const WCHAR *fileName)
{
    int quality = 100; // jpeg only
    wstring wstr = fileName;
    int pos = (int)wstr.rfind(L'.');
    wstring extName = (pos > 0) ? wstr.substr(pos+1) : L"png";
    return image->save(QString::fromWCharArray(fileName), NULL, quality);
}

KinectQtRender::KinectQtRender()
{
    mColorBitmap    = NULL;
    mDepthBitmap    = NULL;
    mDepthFrame     = NULL;
    mSkeletonBitmap = NULL;

    memset(&mSkeletonFrame, 0, sizeof(mSkeletonFrame));

    mSkeletonBackColor.setRgb(255, 255, 255);
    mTrackedLineColor.setRgb(0, 255, 0);
    mInferredLineColor.setRgb(128, 128, 128);
    mTrackedEllipseColor.setRgb(0, 0, 255);
    mInferredEllipseColor.setRgb(128, 128, 128);

    setScaleRatio(1.0f);

    mDrawPen = new QPen(QColor());
    mDrawPen->setCapStyle(Qt::PenCapStyle::RoundCap);
    mDrawPen->setJoinStyle(Qt::PenJoinStyle::RoundJoin);

    mDrawBrush = new QBrush(QColor(0, 0, 0));

    mTextFont = new QFont("Arial", 10);
    mTextFont->setBold(true);
    mTextPen = new QPen(QColor(0, 0, 0));
}

KinectQtRender::~KinectQtRender()
{
    if(mColorBitmap != NULL)
    {
        delete mColorBitmap;
        mColorBitmap = NULL;
    }
    if(mDepthBitmap != NULL)
    {
        delete mDepthBitmap;
        mDepthBitmap = NULL;
    }
    if(mDepthFrame != NULL)
    {
        delete mDepthFrame;
        mDepthFrame = NULL;
    }
    if(mSkeletonBitmap != NULL)
    {
        delete mSkeletonBitmap;
        mSkeletonBitmap = NULL;
    }
    if(mDrawPen != NULL)
    {
        delete mDrawPen;
        mDrawPen = NULL;
    }
    if(mDrawBrush != NULL)
    {
        delete mDrawBrush;
        mDrawBrush = NULL;
    }
    if(mTextFont != NULL)
    {
        delete mTextFont;
        mTextFont = NULL;
    }
    if(mTextPen != NULL)
    {
        delete mTextPen;
        mTextPen = NULL;
    }
}

void KinectQtRender::setScaleRatio(float ratio)
{
    mLineWidth          = ratio*LINEWIDTH;
    mHeadDiameter       = ratio*HEADDIAMETER;
    mHandFootDiameter   = ratio*HANDFOOTDIAMETER;
    mJointDiameter      = ratio*JOINTDIAMETER;

    mScaleRatio         = ratio;
}

QImage * KinectQtRender::getColorBitmap()
{
    return mColorBitmap;
}

QImage * KinectQtRender::getDepthBitmap()
{
    return mDepthBitmap;
}

QImage * KinectQtRender::getSkeletonBitmap()
{
    return mSkeletonBitmap;
}

const NUI_SKELETON_FRAME & KinectQtRender::getSkeletonFrame()
{
    return mSkeletonFrame;
}

void KinectQtRender::drawColorToBitmap(KinectLib *pKinect, int width, int height, const KinectImageFrame &frame, bool drawSkeleton)
{
    verifyBitmap(width, height, QImage::Format::Format_RGB32, mColorBitmap);
    if(mColorBitmap == NULL)
    {
        return ;
    }
    if((width == frame.width) && (height == frame.height))
    {
        copyBitsToBitmap(frame.bits, frame.stride, mColorBitmap);
    }
    else
    {
        QImage frameBitmap((const uchar *)frame.bits, frame.width, frame.height, QImage::Format::Format_RGB32);
        drawBitmapToBitmap(mColorBitmap, &frameBitmap);
    }
    if(drawSkeleton)
    {
        QPainter painter(mColorBitmap);
        drawSkeletonToGraphics(painter, pKinect, width, height, mSkeletonFrame);
    }
}

static void verifyValueRange(int &value, int minValue, int maxValue)
{
    if(value < minValue) { value = minValue; }
    if(value > maxValue) { value = maxValue; }
}

static byte convertDepthToColor(int depth, int minDepth, int maxDepth)
{
//    verifyValueRange(depth, minDepth, maxDepth);
//    int color = ((depth-minDepth)*255)/(maxDepth-minDepth);

//    int color = ((depth >= minDepth) && (depth <= maxDepth)) ? (depth % 256) : 0;

    int realDepth = (depth >> 3);
    int color = (255 - ((255*realDepth)/0x1FFF));

    verifyValueRange(color, 0, 255);
    return byte(color);
}

void KinectQtRender::drawDepthToBitmap(KinectLib *pKinect, int width, int height, const KinectDepthFrame &frame, bool drawSkeleton)
{
    verifyBitmap(width, height, QImage::Format::Format_RGB32, mDepthBitmap);
    verifyBitmap(frame.width, frame.height, QImage::Format::Format_RGB32, mDepthFrame);
    if((mDepthBitmap == NULL) || (mDepthFrame == NULL))
    {
        return ;
    }
    const bool isSameSize = ((width == frame.width) && (height == frame.height));

    QImage *depthBitmap = isSameSize ? mDepthBitmap : mDepthFrame;
    const NUI_DEPTH_IMAGE_PIXEL *depths = frame.depths;
    for(int h = 0; h < frame.height; h ++)
    {
        byte *dst = (byte *)depthBitmap->scanLine(h);
        for(int w = 0; w < frame.width; w ++)
        {
            USHORT depth = depths->depth;
            byte color = convertDepthToColor(depth, frame.minDepth, frame.maxDepth);
            // BBGGRRAA
            dst[w*4+0] = color;
            dst[w*4+1] = color;
            dst[w*4+2] = color;
            dst[w*4+3] = 255;
            depths ++;
        }            
    }
    if(!isSameSize)
    {
        drawBitmapToBitmap(mDepthBitmap, mDepthFrame);
    }
    if(drawSkeleton)
    {
        QPainter painter(mDepthBitmap);
        drawSkeletonToGraphics(painter, pKinect, width, height, mSkeletonFrame);
    }
}

void KinectQtRender::setSkeletonFrame(const NUI_SKELETON_FRAME &frame)
{
    mSkeletonFrame = frame;
}

void KinectQtRender::drawSkeletonToGraphics(QPainter &painter, KinectLib *pKinect, int width, int height, 
                                          const NUI_SKELETON_FRAME &frame)
{
    painter.setRenderHint(QPainter::RenderHint::Antialiasing, true);

    for(int s = 0 ; s < NUI_SKELETON_COUNT; s ++)
    {
        const NUI_SKELETON_TRACKING_STATE trackingState = frame.SkeletonData[s].eTrackingState;
        const NUI_SKELETON_DATA & skeletonData = frame.SkeletonData[s];

        if(trackingState == NUI_SKELETON_TRACKED)
        {
            for(int p = 0; p < NUI_SKELETON_POSITION_COUNT; p ++)
            {
                SkeletonPoint result = pKinect->convertSkeletonPosition(skeletonData.SkeletonPositions[p], width, height);
                mSkeletonPoints[p] = QPoint(result.x, result.y);
            }

            // Render Torso
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_HEAD, NUI_SKELETON_POSITION_SHOULDER_CENTER);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SHOULDER_CENTER, NUI_SKELETON_POSITION_SHOULDER_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SHOULDER_CENTER, NUI_SKELETON_POSITION_SHOULDER_RIGHT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SHOULDER_CENTER, NUI_SKELETON_POSITION_SPINE);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SPINE, NUI_SKELETON_POSITION_HIP_CENTER);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_HIP_CENTER, NUI_SKELETON_POSITION_HIP_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_HIP_CENTER, NUI_SKELETON_POSITION_HIP_RIGHT);

            // Left Arm
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SHOULDER_LEFT, NUI_SKELETON_POSITION_ELBOW_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_ELBOW_LEFT, NUI_SKELETON_POSITION_WRIST_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_WRIST_LEFT, NUI_SKELETON_POSITION_HAND_LEFT);

            // Right Arm
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_SHOULDER_RIGHT, NUI_SKELETON_POSITION_ELBOW_RIGHT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_ELBOW_RIGHT, NUI_SKELETON_POSITION_WRIST_RIGHT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_WRIST_RIGHT, NUI_SKELETON_POSITION_HAND_RIGHT);

            // Left Leg
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_HIP_LEFT, NUI_SKELETON_POSITION_KNEE_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_KNEE_LEFT, NUI_SKELETON_POSITION_ANKLE_LEFT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_ANKLE_LEFT, NUI_SKELETON_POSITION_FOOT_LEFT);

            // Right Leg
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_HIP_RIGHT, NUI_SKELETON_POSITION_KNEE_RIGHT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_KNEE_RIGHT, NUI_SKELETON_POSITION_ANKLE_RIGHT);
            drawBone(painter, skeletonData, NUI_SKELETON_POSITION_ANKLE_RIGHT, NUI_SKELETON_POSITION_FOOT_RIGHT);

            for(int p = 0; p < NUI_SKELETON_POSITION_COUNT; p ++)
            {
                if(skeletonData.eSkeletonPositionTrackingState[p] == NUI_SKELETON_POSITION_TRACKED)
                {
                    mDrawBrush->setColor(mTrackedEllipseColor);
                }
                else if(skeletonData.eSkeletonPositionTrackingState[p] == NUI_SKELETON_POSITION_INFERRED)
                {
                    mDrawBrush->setColor(mInferredEllipseColor);
                }
                else
                {
                    continue ;
                }
                float ellipseDiameter = mJointDiameter;
                switch(p)
                {
                case NUI_SKELETON_POSITION_HEAD: 
                    {
                        ellipseDiameter = mHeadDiameter; 
                    }
                    break ;
                case NUI_SKELETON_POSITION_HAND_LEFT: 
                case NUI_SKELETON_POSITION_HAND_RIGHT: 
                case NUI_SKELETON_POSITION_FOOT_LEFT: 
                case NUI_SKELETON_POSITION_FOOT_RIGHT: 
                    {
                        ellipseDiameter = mHandFootDiameter; 
                    }
                    break ;
                }
                painter.setPen(Qt::PenStyle::NoPen);
                painter.setBrush(*mDrawBrush);
                painter.drawEllipse(mSkeletonPoints[p].x() - (ellipseDiameter/2),
                    mSkeletonPoints[p].y() - (ellipseDiameter/2), ellipseDiameter, ellipseDiameter);
            }
        }
    }
}

void KinectQtRender::drawBone(QPainter &painter, const NUI_SKELETON_DATA & skeletonData, 
                            const NUI_SKELETON_POSITION_INDEX joint0, const NUI_SKELETON_POSITION_INDEX joint1)
{
    NUI_SKELETON_POSITION_TRACKING_STATE joint0State = skeletonData.eSkeletonPositionTrackingState[joint0];
    NUI_SKELETON_POSITION_TRACKING_STATE joint1State = skeletonData.eSkeletonPositionTrackingState[joint1];
    if((joint0State == NUI_SKELETON_POSITION_NOT_TRACKED) || (joint1State == NUI_SKELETON_POSITION_NOT_TRACKED))
    {
        return ;
    }
    if((joint0State == NUI_SKELETON_POSITION_INFERRED) && (joint1State == NUI_SKELETON_POSITION_INFERRED))
    {
        return ;
    }
    if((joint0State == NUI_SKELETON_POSITION_TRACKED) && (joint1State == NUI_SKELETON_POSITION_TRACKED))
    {
        mDrawPen->setColor(mTrackedLineColor);
        mDrawPen->setWidth(mLineWidth);
    }
    else
    {
        mDrawPen->setColor(mInferredLineColor);
        mDrawPen->setWidth(mLineWidth);
    }
    painter.setPen(*mDrawPen);
    painter.drawLine(mSkeletonPoints[joint0], mSkeletonPoints[joint1]);
}

void KinectQtRender::drawSkeletonToBitmap(KinectLib *pKinect, int width, int height, const NUI_SKELETON_FRAME &frame)
{
    verifyBitmap(width, height, QImage::Format::Format_RGB32, mSkeletonBitmap);
    if(mSkeletonBitmap == NULL)
    {
        return ;
    }
    QPainter painter(mSkeletonBitmap);
    painter.fillRect(QRect(0, 0, mSkeletonBitmap->width(), mSkeletonBitmap->height()), QBrush(mSkeletonBackColor));
    drawSkeletonToGraphics(painter, pKinect, width, height, frame);
}

void KinectQtRender::drawPositionTextToGraphics(QPainter &painter, KinectLib *pKinect, int width, int height,
                                              const NUI_SKELETON_FRAME &frame)
{
    painter.setRenderHint(QPainter::RenderHint::TextAntialiasing, true);

    for(int s = 0 ; s < NUI_SKELETON_COUNT; s ++)
    {
        const NUI_SKELETON_TRACKING_STATE trackingState = frame.SkeletonData[s].eTrackingState;
        const NUI_SKELETON_DATA & skeletonData = frame.SkeletonData[s];

        if(trackingState == NUI_SKELETON_TRACKED)
        {
            for(int p = 0; p < NUI_SKELETON_POSITION_COUNT; p ++)
            {
                mSkeleton3Points[p] = pKinect->convertSkeletonPosition(skeletonData.SkeletonPositions[p], width, height);

                if(skeletonData.eSkeletonPositionTrackingState[p] == NUI_SKELETON_POSITION_TRACKED)
                {
                }
                else if(skeletonData.eSkeletonPositionTrackingState[p] == NUI_SKELETON_POSITION_INFERRED)
                {
                }
                else
                {
                    continue ;
                }
                
                QPoint textPoint(mSkeleton3Points[p].x, mSkeleton3Points[p].y);

                wstring wstr = Utils::formatString(L"%d,%d,%d",
                    mSkeleton3Points[p].x, mSkeleton3Points[p].y, mSkeleton3Points[p].depth);

                painter.setFont(*mTextFont);
                painter.setPen(*mTextPen);
                painter.drawText(textPoint, QString::fromStdWString(wstr));
            }
        }
    }
}

void KinectQtRender::drawPositionTextToBitmap(KinectLib *pKinect, int width, int height, const NUI_SKELETON_FRAME &frame)
{
    verifyBitmap(width, height, QImage::Format::Format_RGB32, mSkeletonBitmap);
    if(mSkeletonBitmap == NULL)
    {
        return ;
    }
    QPainter painter(mSkeletonBitmap);
    drawPositionTextToGraphics(painter, pKinect, width, height, frame);
}

void KinectQtRender::drawAngle(QPainter &painter, const NUI_SKELETON_DATA & skeletonData, 
                             const NUI_SKELETON_POSITION_INDEX joint0, 
                             const NUI_SKELETON_POSITION_INDEX joint1, 
                             const NUI_SKELETON_POSITION_INDEX joint2)
{
    NUI_SKELETON_POSITION_TRACKING_STATE joint0State = skeletonData.eSkeletonPositionTrackingState[joint0];
    NUI_SKELETON_POSITION_TRACKING_STATE joint1State = skeletonData.eSkeletonPositionTrackingState[joint1];
    NUI_SKELETON_POSITION_TRACKING_STATE joint2State = skeletonData.eSkeletonPositionTrackingState[joint2];
    if((joint0State == NUI_SKELETON_POSITION_NOT_TRACKED) ||
        (joint1State == NUI_SKELETON_POSITION_NOT_TRACKED) ||
        (joint2State == NUI_SKELETON_POSITION_NOT_TRACKED))
    {
        return ;
    }
    double dAngle = Utils::calcAngle(
        mSkeleton3Points[joint0].x, mSkeleton3Points[joint0].y,
        mSkeleton3Points[joint1].x, mSkeleton3Points[joint1].y, 
        mSkeleton3Points[joint2].x, mSkeleton3Points[joint2].y,
        true);
    QPoint textPoint(mSkeleton3Points[joint1].x, mSkeleton3Points[joint1].y+12);

    wstring wstr = Utils::formatString(L"%.3f", dAngle);

    painter.setFont(*mTextFont);
    painter.setPen(*mTextPen);
    painter.drawText(textPoint, QString::fromStdWString(wstr));
}

void KinectQtRender::drawAngleTextToGraphics(QPainter &painter, KinectLib *pKinect, int width, int height,
                                           const NUI_SKELETON_FRAME &frame)
{
    painter.setRenderHint(QPainter::RenderHint::TextAntialiasing, true);

    for(int s = 0 ; s < NUI_SKELETON_COUNT; s ++)
    {
        const NUI_SKELETON_TRACKING_STATE trackingState = frame.SkeletonData[s].eTrackingState;
        const NUI_SKELETON_DATA & skeletonData = frame.SkeletonData[s];

        if(trackingState == NUI_SKELETON_TRACKED)
        {
            for(int p = 0; p < NUI_SKELETON_POSITION_COUNT; p ++)
            {
                mSkeleton3Points[p] = pKinect->convertSkeletonPosition(skeletonData.SkeletonPositions[p], width, height);
            }
            
            drawAngle(painter, skeletonData, 
                NUI_SKELETON_POSITION_SHOULDER_CENTER, 
                NUI_SKELETON_POSITION_SHOULDER_LEFT,
                NUI_SKELETON_POSITION_ELBOW_LEFT);
            
            drawAngle(painter, skeletonData, 
                NUI_SKELETON_POSITION_SHOULDER_CENTER, 
                NUI_SKELETON_POSITION_SHOULDER_RIGHT,
                NUI_SKELETON_POSITION_ELBOW_RIGHT);

            drawAngle(painter, skeletonData, 
                NUI_SKELETON_POSITION_SHOULDER_LEFT, 
                NUI_SKELETON_POSITION_ELBOW_LEFT,
                NUI_SKELETON_POSITION_WRIST_LEFT);

            drawAngle(painter, skeletonData, 
                NUI_SKELETON_POSITION_SHOULDER_RIGHT, 
                NUI_SKELETON_POSITION_ELBOW_RIGHT,
                NUI_SKELETON_POSITION_WRIST_RIGHT);
        }
    }
}

void KinectQtRender::drawAngleTextToBitmap(KinectLib *pKinect, int width, int height, const NUI_SKELETON_FRAME &frame)
{
    verifyBitmap(width, height, QImage::Format::Format_RGB32, mSkeletonBitmap);
    if(mSkeletonBitmap == NULL)
    {
        return ;
    }
    QPainter painter(mSkeletonBitmap);
    drawAngleTextToGraphics(painter, pKinect, width, height, frame);
}
